﻿#include "mainwindow.h"
#include "./ui_mainwindow.h"

#include <QScreen>
#include <QFileDialog>
#include <QProcess>
#include <QProgressDialog>
#include <QDebug>

#include "forms/forminfo.h"

#include "ogrsf_frmts.h"
#include "cpl_conv.h"
#include "gdal_alg.h"

MainWindow::MainWindow(QWidget *parent)
    : QMainWindow(parent)
    , ui(new Ui::MainWindow)
{
    ui->setupUi(this);

    // 不允许调整窗口大小
    setWindowFlags(windowFlags()& ~Qt::WindowMaximizeButtonHint);
    setFixedSize(this->width(), this->height());
    // 窗口显示在屏幕正中
    QScreen *screen = QGuiApplication::primaryScreen();
    move((screen->geometry().width() - this->width()) / 2, (screen->geometry().height() - this->height()) / 2);

    GDALAllRegister();
    blockSize = 2048;
}

MainWindow::~MainWindow()
{
    delete ui;
}

void MainWindow::on_pushButtonDem_clicked()
{
    ui->lineEditDem->setText(QFileDialog::getOpenFileName(this, QString(), QString(), "GeoTiff(*.tif *.tiff)"));
}

void MainWindow::on_pushButtonClass_clicked()
{
    ui->lineEditClass->setText(QFileDialog::getOpenFileName(this, QString(), QString(), "GeoTiff(*.tif *.tiff)"));
}

void MainWindow::on_pushButtonHabitat_clicked()
{
    QString gdbPath = QFileDialog::getExistingDirectory();
    ui->lineEditHabitat->setText(gdbPath);
    if (gdbPath.isEmpty())
    {
        ui->lineEditHabitat->setText("");
        ui->comboBoxLayer->clear();
        return;
    }

    // check gdb
    GDALDataset *gdbDs = static_cast<GDALDataset*>(GDALOpenEx(gdbPath.toStdString().c_str(), GDAL_OF_VECTOR, NULL, NULL, NULL ));
    if(gdbDs == nullptr)
    {
        msgBox.critical(this, QStringLiteral("文件错误"), QStringLiteral("请选择正确的gdb路径"));
        ui->lineEditHabitat->setText("");
        ui->comboBoxLayer->clear();
        return;
    }

    for (int i = 0; i < gdbDs->GetLayerCount(); i++)
    {
        OGRLayer *layer = gdbDs->GetLayer(i);
        QString name(layer->GetName());
        ui->comboBoxLayer->addItem(name);
    }
    GDALClose(gdbDs);
}

void MainWindow::on_pushButtonResult_clicked()
{
    ui->lineEditResult->setText(QFileDialog::getSaveFileName(this, QString(), QString(), "GeoTiff(*.tif *.tiff)"));
}


void MainWindow::on_pushButtonInfo_clicked()
{
    FormInfo *info = new FormInfo;
    info->show();
}

void MainWindow::on_pushButtonOk_clicked()
{
    // 更新控件
    ui->pushButtonOk->setDisabled(true);

    // 检验文件
    QString demPath = ui->lineEditDem->text();
    QString classPath = ui->lineEditClass->text();
    QString habitatPath = ui->lineEditHabitat->text();
    QString resultPath = ui->lineEditResult->text();
    if (demPath.isEmpty())
    {
        msgBox.critical(this, QStringLiteral("路径错误"), QStringLiteral("请选择DEM文件"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    if (classPath.isEmpty())
    {
        msgBox.critical(this, QStringLiteral("路径错误"), QStringLiteral("请选择分类文件"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    if (habitatPath.isEmpty())
    {
        msgBox.critical(this, QStringLiteral("路径错误"), QStringLiteral("请选择栖息地文件"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    if (resultPath.isEmpty())
    {
        msgBox.critical(this, QStringLiteral("路径错误"), QStringLiteral("请选择结果文件"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }

    // open dem
    GDALDataset *demDs = (GDALDataset*)GDALOpen(demPath.toStdString().c_str(), GA_ReadOnly);
    if (demDs == NULL)
    {
        msgBox.critical(this, QStringLiteral("文件错误"), QStringLiteral("无法打开DEM文件"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    double geoTrans[6];
    demDs->GetGeoTransform(geoTrans);
    const char *proj = demDs->GetProjectionRef();
    GDALRasterBand *demBand = demDs->GetRasterBand(1);
    int xSize = demBand->GetXSize();
    int ySize = demBand->GetYSize();
    double xMin = geoTrans[0];
    double xMax = geoTrans[0] + geoTrans[1] * xSize;
    double yMin = geoTrans[3] + geoTrans[5] * ySize;
    double yMax = geoTrans[3];

    // open class
    GDALDataset *classDs = (GDALDataset*)GDALOpen(classPath.toStdString().c_str(), GA_ReadOnly);
    if (classDs == NULL)
    {
        GDALClose((GDALDatasetH) demDs);
        msgBox.critical(this, QStringLiteral("文件错误"), QStringLiteral("无法打开分类文件"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    GDALRasterBand *classBand = classDs->GetRasterBand(1);
    if (classBand->GetXSize() != xSize || classBand->GetYSize() != ySize)
    {
        GDALClose((GDALDatasetH) classDs);
        GDALClose((GDALDatasetH) demDs);
        msgBox.critical(this, QStringLiteral("文件错误"), QStringLiteral("DEM文件与分类文件行列数不匹配"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }

    // open habitat
    GDALDataset *habitatDs = static_cast<GDALDataset*>(GDALOpenEx(habitatPath.toStdString().c_str(), GDAL_OF_VECTOR, NULL, NULL, NULL));
    if (habitatDs == NULL)
    {
        GDALClose((GDALDatasetH) demDs);
        GDALClose((GDALDatasetH) classDs);
        msgBox.critical(this, QStringLiteral("文件错误"), QStringLiteral("无法打开栖息地文件"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    OGRLayer *habitatLayer = habitatDs->GetLayerByName(ui->comboBoxLayer->currentText().toStdString().c_str());

    // check fields
    OGRFeatureDefn *defn = habitatLayer->GetLayerDefn();
    if (defn->GetFieldIndex("altitude_low") < 0)
    {
        GDALClose((GDALDatasetH) demDs);
        GDALClose((GDALDatasetH) classDs);
        GDALClose(habitatDs);
        msgBox.critical(this, QStringLiteral("数据错误"), QStringLiteral("所选图层未包含altitude_low字段"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    if (defn->GetFieldIndex("altitude_upper") < 0)
    {
        GDALClose((GDALDatasetH) demDs);
        GDALClose((GDALDatasetH) classDs);
        GDALClose(habitatDs);
        msgBox.critical(this, QStringLiteral("数据错误"), QStringLiteral("所选图层未包含altitude_upper字段"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    if (defn->GetFieldIndex("h1") < 0)
    {
        GDALClose((GDALDatasetH) demDs);
        GDALClose((GDALDatasetH) classDs);
        GDALClose(habitatDs);
        msgBox.critical(this, QStringLiteral("数据错误"), QStringLiteral("所选图层未包含h1字段"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    if (defn->GetFieldIndex("h2") < 0)
    {
        GDALClose((GDALDatasetH) demDs);
        GDALClose((GDALDatasetH) classDs);
        GDALClose(habitatDs);
        msgBox.critical(this, QStringLiteral("数据错误"), QStringLiteral("所选图层未包含h2字段"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    if (defn->GetFieldIndex("h3") < 0)
    {
        GDALClose((GDALDatasetH) demDs);
        GDALClose((GDALDatasetH) classDs);
        GDALClose(habitatDs);
        msgBox.critical(this, QStringLiteral("数据错误"), QStringLiteral("所选图层未包含h3字段"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    if (defn->GetFieldIndex("h4") < 0)
    {
        GDALClose((GDALDatasetH) demDs);
        GDALClose((GDALDatasetH) classDs);
        GDALClose(habitatDs);
        msgBox.critical(this, QStringLiteral("数据错误"), QStringLiteral("所选图层未包含h4字段"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    if (defn->GetFieldIndex("h5") < 0)
    {
        GDALClose((GDALDatasetH) demDs);
        GDALClose((GDALDatasetH) classDs);
        GDALClose(habitatDs);
        msgBox.critical(this, QStringLiteral("数据错误"), QStringLiteral("所选图层未包含h5字段"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    if (defn->GetFieldIndex("h6") < 0)
    {
        GDALClose((GDALDatasetH) demDs);
        GDALClose((GDALDatasetH) classDs);
        GDALClose(habitatDs);
        msgBox.critical(this, QStringLiteral("数据错误"), QStringLiteral("所选图层未包含h6字段"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    if (defn->GetFieldIndex("h7") < 0)
    {
        GDALClose((GDALDatasetH) demDs);
        GDALClose((GDALDatasetH) classDs);
        GDALClose(habitatDs);
        msgBox.critical(this, QStringLiteral("数据错误"), QStringLiteral("所选图层未包含h7字段"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }
    if (defn->GetFieldIndex("h8") < 0)
    {
        GDALClose((GDALDatasetH) demDs);
        GDALClose((GDALDatasetH) classDs);
        GDALClose(habitatDs);
        msgBox.critical(this, QStringLiteral("数据错误"), QStringLiteral("所选图层未包含h8字段"));
        ui->pushButtonOk->setDisabled(false);
        return;
    }

    // init progress
    int featureCount = habitatLayer->GetFeatureCount();
    QProgressDialog proDialog(QStringLiteral("初始化结果文件"), QStringLiteral("取消"), 0, featureCount, this);
    proDialog.setModal(true);
    proDialog.setWindowTitle(QStringLiteral("计算进度"));
    proDialog.setValue(0);
    proDialog.show();
    QCoreApplication::processEvents();
    if (proDialog.wasCanceled())
    {
        GDALClose((GDALDatasetH) demDs);
        GDALClose((GDALDatasetH) classDs);
        GDALClose(habitatDs);
        ui->pushButtonOk->setDisabled(false);
        return;
    }

    // create habitat
    GDALDriver *outDriver = GetGDALDriverManager()->GetDriverByName("GTiff");
    char **papszOptions = NULL;
    papszOptions = CSLSetNameValue(papszOptions, "TILED", "YES");
    papszOptions = CSLSetNameValue(papszOptions, "COMPRESS", "LZW");
    GDALDataset *resultDs = outDriver->Create(resultPath.toStdString().c_str(), xSize, ySize, 1, GDT_UInt16, papszOptions);
    resultDs->SetGeoTransform(geoTrans);
    resultDs->SetProjection(proj);
    GDALRasterBand *resultBand = resultDs->GetRasterBand(1);
    resultBand->SetNoDataValue(65535);
    CSLDestroy(papszOptions);

    // init result raster
    initResult(demBand, classBand, resultBand);

    // scan all feature
    int calculated = 0;
    proDialog.setLabelText(QStringLiteral("正在计算：已完成 %1 / %2").arg(0).arg(featureCount));
    habitatLayer->ResetReading();
    OGRFeature *feature;
    while ((feature = habitatLayer->GetNextFeature()) != NULL) // per feature
    {
        // get extent
        OGRGeometry *geometry = feature->GetGeometryRef();
        OGREnvelope *envelope = new OGREnvelope();
        geometry->getEnvelope(envelope);
        double gMinX = envelope->MinX;
        double gMinY = envelope->MinY;
        double gMaxX = envelope->MaxX;
        double gMaxY = envelope->MaxY;

        // fix extent to fit raster pixel
        int xOffset = 0;
        int yOffset = 0;
        if (gMinX < xMin)
        {
            gMinX = xMin;
        }
        else
        {
            xOffset = (gMinX - geoTrans[0]) / geoTrans[1];
            gMinX = geoTrans[0] + geoTrans[1] * xOffset;
        }
        if (gMaxX > xMax)
        {
            gMaxX = xMax;
        }
        if (gMinY < yMin)
        {
            gMinY = yMin;
        }
        if (gMaxY > yMax)
        {
            gMaxY = yMax;
        }
        else
        {
            yOffset = (gMaxY - geoTrans[3]) / geoTrans[5];
            gMaxY = geoTrans[3] + geoTrans[5] * yOffset;
        }
        int tmpXSize = (gMaxX - gMinX) / geoTrans[1] + 0.5;
        int tmpYSize = (gMinY - gMaxY) / geoTrans[5] + 0.5;
        if (tmpXSize < 1 || tmpYSize < 1) // less than one pixel
        {
            proDialog.setValue(calculated);
            proDialog.setLabelText(QStringLiteral("正在计算：已完成 %1 / %2").arg(++calculated).arg(featureCount));
            QCoreApplication::processEvents();
            OGRFeature::DestroyFeature(feature);
            continue;
        }
        double tmpGeoTrans[6] = { gMinX, geoTrans[1], 0, gMaxY, 0, geoTrans[5] };

        // rasterize
        GDALDataset *tmpDs = rasterize(geometry, tmpXSize, tmpYSize, tmpGeoTrans);
        GDALRasterBand *tmpBand = tmpDs->GetRasterBand(1);

        updateHabitat(demBand, classBand, resultBand, tmpBand, feature, blockSize, xOffset, yOffset);

        // clear feature with raster
        GDALClose((GDALDatasetH)tmpDs);
        OGRFeature::DestroyFeature(feature);

        // update progress dialog
        proDialog.setValue(calculated);
        proDialog.setLabelText(QStringLiteral("正在计算：已完成 %1 / %2").arg(++calculated).arg(featureCount));
        QCoreApplication::processEvents();
        if (proDialog.wasCanceled())
        {
            GDALClose((GDALDatasetH) demDs);
            GDALClose((GDALDatasetH) classDs);
            GDALClose((GDALDatasetH) resultDs);
            GDALClose(habitatDs);
            ui->pushButtonOk->setDisabled(false);
            return;
        }
    }

    GDALClose((GDALDatasetH) demDs);
    GDALClose((GDALDatasetH) classDs);
    GDALClose((GDALDatasetH) resultDs);
    GDALClose(habitatDs);
    proDialog.setValue(featureCount);
    msgBox.information(this, QStringLiteral("计算完成"), QStringLiteral("计算完成"));
    ui->pushButtonOk->setDisabled(false);
}

void MainWindow::on_pushButtonCancel_clicked()
{
    this->close();
}

void MainWindow::initResult(GDALRasterBand *demBand, GDALRasterBand *classBand, GDALRasterBand *resultBand)
{
    int xSize = demBand->GetXSize();
    int ySize = demBand->GetYSize();
    int demNoData = demBand->GetNoDataValue();
    int classNoData = classBand->GetNoDataValue();

    // process data blocked
    int xOffset = 0;
    int yOffset = 0;
    int leftXSize = xSize;
    int leftYSize = ySize;
    while (leftYSize > 0)
    {
        int calYSize = (leftYSize > blockSize) ? blockSize : leftYSize;
        yOffset = ySize - leftYSize;

        while (leftXSize > 0)
        {
            int calXSize = (leftXSize > blockSize) ? blockSize : leftXSize;
            xOffset = xSize - leftXSize;

            // read data
            int calTotalSize = calXSize * calYSize;
            GDALDataType demType = demBand->GetRasterDataType();
            GDALDataType classType = classBand->GetRasterDataType();
            void *demData = initDataBlock(demType, calTotalSize);
            void *classData = initDataBlock(classType, calTotalSize);
            GUInt16 *resultData = (GUInt16*) CPLMalloc(sizeof(GUInt16) * calTotalSize);

            demBand->RasterIO(GF_Read, xOffset, yOffset, calXSize, calYSize, demData, calXSize, calYSize, demType, 0, 0);
            classBand->RasterIO(GF_Read, xOffset, yOffset, calXSize, calYSize, classData, calXSize, calYSize, classType, 0, 0);

            for (int i = 0; i < calTotalSize; i++)
            {
                int demValue = getData(demData, demType, i);
                int classValue = getData(classData, classType, i);

                if (demValue == demNoData || classValue == classNoData)
                {
                    resultData[i] = 65535;
                }
                else
                {
                    resultData[i] = 0;
                }
            }

            resultBand->RasterIO(GF_Write, xOffset, yOffset, calXSize, calYSize, resultData, calXSize, calYSize, GDT_UInt16, 0, 0);

            CPLFree(demData);
            CPLFree(classData);
            CPLFree(resultData);
            leftXSize = leftXSize - calXSize;
        }

        leftYSize = leftYSize - calYSize;
        // reset x
        leftXSize = xSize;
    }
}

void *MainWindow::initDataBlock(GDALDataType type, int size)
{
    switch (type)
    {
    case GDT_Byte:
        return CPLMalloc(sizeof(GByte) * size);
    case GDT_UInt16:
        return CPLMalloc(sizeof(GUInt16) * size);
    case GDT_Int16:
        return CPLMalloc(sizeof(GInt16) * size);
    case GDT_UInt32:
        return CPLMalloc(sizeof(GUInt32) * size);
    default:
        return CPLMalloc(sizeof(int) * size);
    }
}

int MainWindow::getData(void *data, GDALDataType type, int idx)
{
    switch (type)
    {
    case GDT_Byte:
        return ((GByte*) data)[idx];
    case GDT_UInt16:
        return ((GUInt16*) data)[idx];
    case GDT_Int16:
        return ((GInt16*) data)[idx];
    case GDT_UInt32:
        return ((GUInt32*) data)[idx];
    default:
        return ((int*) data)[idx];
    }
}

GDALDataset *MainWindow::rasterize(OGRGeometry *geometry, int xSize, int ySize, double *geoTrans)
{
    GDALDriver *outDriver = GetGDALDriverManager()->GetDriverByName("GTiff");
    char **papszOptions = NULL;
    papszOptions = CSLSetNameValue(papszOptions, "TILED", "YES");
    papszOptions = CSLSetNameValue(papszOptions, "COMPRESS", "LZW");
    GDALDataset *outDs = outDriver->Create("/vsimem/f07af6c5.tif", xSize, ySize, 1, GDT_Byte, papszOptions);
    outDs->SetGeoTransform(geoTrans);
    CSLDestroy(papszOptions);

    // burn data
    double burnValue[1] = { 1.0 };
    int bandList[1] = { 1 };
    OGRGeometryH geoms[1] = { (OGRGeometryH) geometry };

    // gdal rasterize opts
    char** optStrs = NULL;
    optStrs = CSLAddString(optStrs, "-optim");
    optStrs = CSLAddString(optStrs, "VECTOR");

    GDALRasterizeGeometries(outDs, 1, bandList, 1, geoms, NULL, NULL, burnValue, optStrs, NULL, NULL);

    CSLDestroy(optStrs);
    return outDs;
}

void MainWindow::updateHabitat(GDALRasterBand *demBand, GDALRasterBand *classBand, GDALRasterBand *resultBand, GDALRasterBand *tmpBand,
                               OGRFeature *feature, int blockSize, int xOffset, int yOffset)
{
    // start calculate
    int tmpXOffset = 0;
    int tmpYOffset = 0;
    int tmpXSize = tmpBand->GetXSize();
    int tmpYSize = tmpBand->GetYSize();
    int leftXSize = tmpXSize;
    int leftYSize = tmpYSize;
    while (leftYSize > 0)
    {
        int calYSize = (leftYSize > blockSize) ? blockSize : leftYSize;
        tmpYOffset = tmpYSize - leftYSize;

        while (leftXSize > 0)
        {
            int calXSize = (leftXSize > blockSize) ? blockSize : leftXSize;
            tmpXOffset = tmpXSize - leftXSize;

            // read data
            int calTotalSize = calXSize * calYSize;
            GDALDataType demType = demBand->GetRasterDataType();
            GDALDataType classType = classBand->GetRasterDataType();
            void *demData = initDataBlock(demType, calTotalSize);
            void *classData = initDataBlock(classType, calTotalSize);
            GUInt16 *resultData = (GUInt16*) CPLMalloc(sizeof(GUInt16) * calTotalSize);
            GByte *tmpData = (GByte*)CPLMalloc(sizeof(GByte) * calTotalSize);

            int imgXOffset = xOffset + tmpXOffset;
            int imgYOffset = yOffset + tmpYOffset;
            demBand->RasterIO(GF_Read, imgXOffset, imgYOffset, calXSize, calYSize, demData, calXSize, calYSize, demType, 0, 0);
            classBand->RasterIO(GF_Read, imgXOffset, imgYOffset, calXSize, calYSize, classData, calXSize, calYSize, classType, 0, 0);
            resultBand->RasterIO(GF_Read, imgXOffset, imgYOffset, calXSize, calYSize, resultData, calXSize, calYSize, GDT_UInt16, 0, 0);
            tmpBand->RasterIO(GF_Read, tmpXOffset, tmpYOffset, calXSize, calYSize, tmpData, calXSize, calYSize, GDT_Byte, 0, 0);

            for (int i = 0; i < calTotalSize; i++)
            {
                int demValue = getData(demData, demType, i);
                int classValue = getData(classData, classType, i);
                GUInt16 resultValue = resultData[i];
                GByte tmpValue = tmpData[i];

                if (resultValue == 65535 || tmpValue == 0)
                {
                    continue;
                }

                resultData[i] = resultValue + checkHabitat(feature, demValue, classValue);
            }

            resultBand->RasterIO(GF_Write, imgXOffset, imgYOffset, calXSize, calYSize, resultData, calXSize, calYSize, GDT_UInt16, 0, 0);

            CPLFree(demData);
            CPLFree(classData);
            CPLFree(resultData);
            CPLFree(tmpData);
            leftXSize = leftXSize - calXSize;
        }

        leftYSize = leftYSize - calYSize;
        // reset x
        leftXSize = tmpXSize;
    }
}

int MainWindow::checkHabitat(OGRFeature* feature, int demValue, int classValue)
{
    if (!feature->IsFieldNull(feature->GetFieldIndex("altitude_low"))
            && !feature->IsFieldNull(feature->GetFieldIndex("altitude_upper")))
    {
        if (demValue < feature->GetFieldAsInteger("altitude_low") || demValue > feature->GetFieldAsInteger("altitude_upper"))
        {
            return 0;
        }
    }

    // check land
    if (!feature->IsFieldNull(feature->GetFieldIndex("h1")))
    {
        if (classValue == feature->GetFieldAsInteger("h1"))
        {
            return 1;
        }
    }
    if (!feature->IsFieldNull(feature->GetFieldIndex("h2")))
    {
        if (classValue == feature->GetFieldAsInteger("h2"))
        {
            return 1;
        }
    }
    if (!feature->IsFieldNull(feature->GetFieldIndex("h3")))
    {
        if (classValue == feature->GetFieldAsInteger("h3"))
        {
            return 1;
        }
    }
    if (!feature->IsFieldNull(feature->GetFieldIndex("h4")))
    {
        if (classValue == feature->GetFieldAsInteger("h4"))
        {
            return 1;
        }
    }
    if (!feature->IsFieldNull(feature->GetFieldIndex("h5")))
    {
        if (classValue == feature->GetFieldAsInteger("h5"))
        {
            return 1;
        }
    }
    if (!feature->IsFieldNull(feature->GetFieldIndex("h6")))
    {
        if (classValue == feature->GetFieldAsInteger("h6"))
        {
            return 1;
        }
    }
    if (!feature->IsFieldNull(feature->GetFieldIndex("h7")))
    {
        if (classValue == feature->GetFieldAsInteger("h7"))
        {
            return 1;
        }
    }
    if (!feature->IsFieldNull(feature->GetFieldIndex("h8")))
    {
        if (classValue == feature->GetFieldAsInteger("h8"))
        {
            return 1;
        }
    }

    return 0;
}
